import subprocess
subprocess.run(['pip', 'uninstall', '-y', 'apex'], check=True)
import gradio as gr
import numpy as np
import random
import torch
from diffusers import StableDiffusion3Pipeline
import qwen_api





device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16
repo = "hf-models/stable-diffusion-3.5-large"


pipe = StableDiffusion3Pipeline.from_pretrained(
    repo, torch_dtype=dtype, device_map="balanced") #.to("cuda") # balanced
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024



def infer(prompt, negative_prompt="Disabled feet, abnormal feet, deformed, distorted, disfigured, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, NSFW", seed=1, randomize_seed=True, width=1024, height=1024, guidance_scale=7, num_inference_steps=32, use_zh=True, progress=gr.Progress(track_tqdm=True)):
    print("使用中文", use_zh)
    resPrompt = ""
    if (use_zh):
        try:
            qwenRes = qwen_api.query({
                "messages": [
                    {
                        "role": "system",
                        "content": "你是一个文生图 Prompt 翻译器，将中文翻译为英文，翻译简洁精准，如果已经是英文就直接原文返回。你只需要直接给出英文答案, 不要废话"
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                "stream": False

            })
            resPrompt = qwenRes["choices"][0]["message"]["content"]
        except Exception as e:
            print("请求失败")
            print(f"Error: {e}")
            resPrompt = prompt
            raise gr.Error("解析中文提示词失败，请稍后或取消勾选中文提示词")
    else:
        resPrompt = prompt

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    print(prompt)
    print(resPrompt)

    generator = torch.Generator().manual_seed(seed)

    image = pipe(
        prompt=resPrompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width, # image_processor.py:111: RuntimeWarning: invalid value encountered in cast
        height=width,
        generator=generator
    ).images[0]

    torch.cuda.empty_cache()

    return image, seed


examples = [
    "美丽的白色樱花盛开，在晴朗的天空背景下，树枝上精致而新鲜的花朵，柔和对焦的摄影，日式风格，微距镜头，自然光，春天的气息，色调清新",
    "云销雨霁，彩彻区明，落霞与孤鹜齐飞，秋水共长天一色",
    "木框，冰雕，极其细致和复杂，华丽，超现实主义，光滑，焦点清晰，绘画，数字艺术，概念艺术，插图，艺术趋势，象征性和有意义的风格，8K",
    "一幅异想天开且富有创意的图像，描绘了一种华夫饼和河马的混合生物，在早餐主题景观中沐浴在融化的黄油河中。它具有河马独特、笨重的体形。然而，这种生物的身体并不是通常的灰色皮肤，而是像刚从煎锅上出来的金棕色酥脆华夫饼。皮肤上有熟悉的华夫饼网格图案，每个方块都充满了闪闪发光的糖浆。环境将河马的自然栖息地与早餐桌环境的元素结合在一起，一条温暖融化的黄油河，背景中郁郁葱葱的煎饼状树叶中露出超大的器皿或盘子，一座高耸的胡椒磨坊矗立在对于一棵树。 当太阳在这个奇幻的世界中升起时，它会给整个场景带来温暖、黄油般的光芒。这个生物在黄油河中感到满足，打了个哈欠。附近，一群鸟儿飞翔",

]
css = """
#col-container {
    margin: 0 auto;
    max-width: 580px;
}
"""
with gr.Blocks(css=css) as demo:

    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""
        # stable-diffusion-3.5-large
        - 支持使用中文提示词, 由 [Gitee AI Serverless API](https://ai.gitee.com/serverless-api) 支持
        - [前往使用文生图 Prompt 辅助工具](https://ai.gitee.com/apps/stringify/image-prompt-tool)
        """)

        with gr.Row():

            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                lines=3,
                placeholder="输入你的提示词",
                container=False,
            )

            run_button = gr.Button("运行", scale=0)

        with gr.Row():
            use_zh = gr.Checkbox(label="使用中文提示词", value=True)

        result = gr.Image(label="Result", show_label=False)
        with gr.Accordion("环境配置", open=False):

            negative_prompt = gr.Text(
                label="否定提示词",
                value="Disabled feet, abnormal feet, deformed, distorted, disfigured, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, NSFW",
                max_lines=1,
                placeholder="输入否定提示词",
            )

            seed = gr.Slider(
                label="种子",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )

            randomize_seed = gr.Checkbox(label="使用随机种子", value=True)

            with gr.Row():

                width = gr.Slider(
                    label="宽度",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=1024,
                )

                height = gr.Slider(
                    label="高度",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=1024,
                )

            with gr.Row():

                guidance_scale = gr.Slider(
                    label="引导比例",
                    minimum=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=7.0,
                )

                num_inference_steps = gr.Slider(
                    label="推理步骤数",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )

        gr.Examples(
            label="例子",
            examples=examples,
            cache_examples=True,
            outputs=[result, seed],
            fn=infer,
            inputs=[prompt],
        )
    gr.on(
        triggers=[run_button.click, prompt.submit, negative_prompt.submit],
        fn=infer,
        inputs=[prompt, negative_prompt, seed, randomize_seed,
                width, height, guidance_scale, num_inference_steps, use_zh],
        outputs=[result, seed]
    )

demo.queue("auto")
demo.launch(show_api=False)
